
from typing import List, Dict
import os
import dashscope


def qwen_call(model="qwen3-max", convs=None):
    if convs is None:
        convs = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "中国的首都是是哪？简短的告诉我。"}
        ]

    response = dashscope.Generation.call(
        # 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：api_key="sk-xxx"
        api_key=os.getenv('DASHSCOPE_API_KEY'),
        # 此处以qwen-plus-2025-04-28为例，可按需更换为支持联网搜索的模型
        model=model,
        messages=convs,
        enable_thinking=True, # 开启深度思考的参数，对qwen3-max似乎无效
        stream=True,
        incremental_output=True,
        result_format="message",
    )

    # 定义完整思考过程
    reasoning_content = ""
    # 定义完整回复
    answer_content = ""
    # 判断是否结束思考过程并开始回复
    is_answering = False
    # 判断是否为第一个chunk，便于打印搜索信息
    is_first_chunk = True

    # print("=" * 20 + "搜索信息" + "=" * 20)

    search_results = []

    for chunk in response:
        if is_first_chunk:
            if "search_info" in chunk.output:
                search_results = chunk.output.search_info["search_results"]
            # for web in search_results:
            #     print(f"[{web['index']}]: [{web['title']}]({web['url']})")
            # print("=" * 20 + "思考过程" + "=" * 20)
            is_first_chunk = False

        # 如果思考过程与回复皆为空，则忽略
        if (chunk.output.choices[0].message.content == "" and
                "reasoning_content" in chunk.output.choices[0].message and
                chunk.output.choices[0].message.reasoning_content == ""):
            pass
        else:
            # 如果当前为思考过程
            if ("reasoning_content" in chunk.output.choices[0].message and
                    chunk.output.choices[0].message.reasoning_content != "" and
                    chunk.output.choices[0].message.content == ""):
                # print(chunk.output.choices[0].message.reasoning_content, end="", flush=True)
                reasoning_content += chunk.output.choices[0].message.reasoning_content
            # 如果当前为回复
            elif chunk.output.choices[0].message.content != "":
                if not is_answering:
                    # print("\n" + "=" * 20 + "完整回复" + "=" * 20)
                    is_answering = True
                # print(chunk.output.choices[0].message.content, end="", flush=True)
                answer_content += chunk.output.choices[0].message.content

    # 如果您需要打印完整思考过程与完整回复，请将以下代码解除注释后运行
    # print("=" * 20 + "完整思考过程" + "=" * 20 + "\n")
    # print(f"{reasoning_content}")
    # print("=" * 20 + "完整回复" + "=" * 20 + "\n")
    # print(f"{answer_content}")
    # 如果您需要打印本次请求的 Token 消耗，请将以下代码解除注释后运行
    # print("\n"+"="*20+"Token 消耗"+"="*20)
    # print(chunk.usage)
    return {"search": search_results, "thinking": reasoning_content, "answer": answer_content, "usage": chunk.usage}


if __name__ == '__main__':
    res = qwen_call()
    print(res)
